package com.ai.module.websocket.server;

import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.constant.RedisKeyStringConstant;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.chat.common.enu.ModelType;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.websocket.coder.UserMessageDecoder;
import com.ai.module.websocket.coder.UserMessageEncoder;
import com.ai.module.business.domain.vo.UserMessage;
import com.ai.module.user.domain.vo.UserLoginInfo;
import com.ai.module.user.service.ChatService;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.RedisTemplateUtils;
import com.ai.module.chat.util.SpringUtils;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiModelProperty;
import lombok.extern.slf4j.Slf4j;
import org.junit.platform.commons.util.StringUtils;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: MrHe
 * @create: 2024-01-27 19:23
 * @description:
 **/

@Component
@ServerEndpoint(value = "/webSocket/{robotId}",
    encoders = UserMessageEncoder.class,
    decoders = UserMessageDecoder.class)
@Slf4j
@Api(tags = "用户websocket聊天 路径参数 AIid")
public class WebSocketController {

    private static ConcurrentHashMap<String,Session> webSocketMap = new ConcurrentHashMap<>();
    private Session session;


    @OnOpen
    @ApiModelProperty("开始连接")
    public void onOpen(Session session, @PathParam("robotId") String robotId) throws EncodeException, IOException {
        this.session = session;
        webSocketMap.put(robotId,session);
        log.info("原始 websocket 有新连接加入！当前在线人数为 [{}] ",webSocketMap.keySet().size());
        UserMessage userMessage = new UserMessage();
        userMessage.setUserMessage(StringPool.EMPTY);
        // sendMessage(session,userMessage);
        // 进行逻辑处理
        // session.getBasicRemote().sendObject(userMessage);
        log.info("用户[{}]加入链接[{}]",ThreadLocalContext.getUserId(),robotId);
        //
    }

    @OnClose
    public void onClose() {
        webSocketMap = new ConcurrentHashMap<>();
        log.info("有一连接关闭！" );
    }

    @OnMessage
    @ApiModelProperty("接受数据")
    public void onMessage( UserMessage message, Session session) {

        log.info("来自客户端的消息 object[{}]", message);
    }
    @OnMessage
    public void onMessage(String message, Session session) {
        if(message.equals("ping")){
            log.info("接收到 websocket 断线重连信息 [{}]", DateUtils.getCurrentTime());
        }else {
            UserMessage userMessage = JSONObject.parseObject(message, UserMessage.class);
            String realMessage = userMessage.getUserMessage();
            if (StringUtils.isBlank(realMessage) || StringUtils.isBlank(realMessage.trim())) {
                log.error("用户对话收到数据为空[{}]", realMessage);
                userMessage.setUserMessage("请输入您的问题");
                sendMessage(session, userMessage);
                return;
            }
            String loginToken = userMessage.getLoginToken();
            dealToken(loginToken);
            ModelType modelType = userMessage.getModelType();
            ChatService service = SpringUtils.getClazz(ChatService.class);
            if (modelType == ModelType.DALLA) {
                List<String> image = service.createImage(realMessage);
                UserMessage aiMessage = new UserMessage();
                aiMessage.setEnd(true);
                aiMessage.setImageList(image);
                sendMessage(session, aiMessage);
            } else {
                service.talkWebSocket(userMessage, this, session);

            }
        }

    }

    @OnError
    public void onError(Session session, Throwable error) {
        System.out.println("发生错误");
        error.printStackTrace();
    }

    public void sendMessage(Session session, UserMessage userMessage) {
        try {
            session.getBasicRemote().sendObject(userMessage);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }



    /**
     * 发送自定义消息
     */
//    public static void sendInfo(String message, String robotId) {
//        try {
//            if (webSocketMap.containsKey(robotId)) {
//                Session session = webSocketMap.get(robotId);
//                if (session.getBasicRemote() != null) {
//                    session.getBasicRemote().sendText(message);
//                }
//            } else {
//                log.error("订阅类型：" + robotId + ",不存在！");
//                throw new Exception("连接已关闭，请刷新页面后重试");
//            }
//        }catch (Exception ex){
//            ex.printStackTrace();
//        }
//
//
//    }
    static void dealToken(String headToken) {
        if (StringUtils.isBlank(headToken)){
            throw new BusinessException( 401,"请重新登录");
        }
        String userTokenKey = String.format(RedisKeyStringConstant.USER_INFO, headToken);

        UserLoginInfo userLoginInfo =  RedisTemplateUtils.get(userTokenKey);
        if (userLoginInfo ==null){
            throw new BusinessException("登录已过期,请重新登录",401);
        }
        ThreadLocalContext.set(StringConstant.USER_INFO,userLoginInfo);
        ThreadLocalContext.set(StringConstant.CURRENT_TIME,System.currentTimeMillis());
        ThreadLocalContext.setUserLoginInfo(userLoginInfo);
    }

}
